package constructdata.TestDataConstructors;

import codetree.CodeTree;
import codetree.CodeTreeOperation;
import codetree.TreeNode;
import treeview.DisplayTreeView;
import treeview.TreeView;

import java.io.*;
import java.util.*;
import java.util.stream.Collectors;

/**
 * Construct test data (Preceding Context)
 * Created by wangxin on 2018/02/05.
 */
public class PreTestDataConstructor {

    private boolean isDebug = false;

    // The marks
    private final String[] filterSigns = {"[","]"};
    private final String[] specialLabel = {"end", "conditionEnd"};

    // The construct trees
    private LinkedList<CodeTree> trees;
    // The predication holes correspondent to the constructTrees;
    private LinkedList<String> predications;
    // CodeTreeOperator
    private CodeTreeOperation operator;
    // The predicate node's complete class name
    private LinkedList<String> classnames;
    // The node who is the parent of the predicated one
    private LinkedList<TreeNode> parents;
    // The size of hole
    private  LinkedList<String> holesizes;
    // The block predictions
    private LinkedList<String> blockpredictions;
    private String blockprediction;
    // The original statements
    private LinkedList<String> originalStatements;

    // The block original statements
    private LinkedList<String> blockstatements;
    private String blockstatement;

    // The variable names
    private LinkedList<String> variableNames;
    private final String[] holeParentAdjust = {"else","elseif","catch","finally","case","default"};

    // testcase trace
    private LinkedList<String> testcaseTraces;
    // testcase pool: <trace, testcase>
    private Map<String, String> testcasePool;
    private Map<Integer,String> sourceLines;
    private String trace;
    private String testcaseDir;

    // test
    private int count = 0;
    /**
     * Constructor
     * */
    public PreTestDataConstructor(){
        operator = new CodeTreeOperation();
        this.count  = 0;
    }

    /**
     * Construct training tree from the input code tree
     * Return the LinkedList with 2 list:
     * 1. The training trees with some continuous holes;
     * 2. The predication holes correspondent to the training trees;
     * 3. And save the trees and predicates to the file
     * */
    public void construct(CodeTree tree,
                          FileWriter treeWriter,FileWriter predictionWriter,FileWriter classWriter, FileWriter generationNodeWriter,FileWriter treeSentenceWriter, FileWriter jarWriter, FileWriter holeSizeWriter,
                          FileWriter traceWriter, // trace back
                          FileWriter blockpredictionWriter, // block of predictions (more lines)
                          FileWriter originalStatementsWriter, // original statements
                          FileWriter variableNamesWriter,// variable names
                          FileWriter linesWriter,// lines writer
                          FileWriter testCaseTraceWriter,
                          FileWriter blockOriginalStatementWriter,
                          String testcaseDir,
                          boolean isCompleteFlag){

        this.testcaseDir = testcaseDir;

        LinkedList<LinkedList> result = getConstructTrainingData(tree);
        List<CodeTree> treeList = result.get(0);
        List<String> predictionList = result.get(1);
        List<String> classList = result.get(2);
        List<TreeNode> generationNodeList = result.get(3);
        List<String> holeSizeList = result.get(4);// record size of hole
        List<String> blockpredictionList = result.get(5);
        List<String> originalStatementsList = result.get(6);
        List<String> variableNameList = result.get(7);
        List<String> testCaseTraceList = result.get(8); // -- testcase trace
        List<String> blockOriginalStatementsList = result.get(9);

        //FindJarHandler findJarHandler = new FindJarHandler();
        for (int i = 0; i < treeList.size(); i++) {
            try {
                //String jar = findJarHandler.getPackage(classList.get(i));
                //if(jar != null) {
                CodeTree tempTree = treeList.get(i);

                operator.saveRegularizedTreeInFile(operator.regularization(tempTree), treeWriter);
                operator.saveTrainingPredictionInFile(predictionList.get(i), predictionWriter);
                operator.saveTrainingPredictionInFile(classList.get(i), classWriter);
                //operator.saveTrainingPredictionInFile(jar, jarWriter);
                int parentnum = generationNodeList.get(i).getSerialNumber();
                operator.saveTrainingPredictionInFile(parentnum + " " + ((parentnum != 0) ? generationNodeList.get(i).getCompleteMethodDeclaration() : ""), generationNodeWriter);
                operator.saveTreeStringFormatInFile(tempTree, treeSentenceWriter, isCompleteFlag);

                operator.saveTrainingPredictionInFile(holeSizeList.get(i), holeSizeWriter);

                operator.saveTrainingPredictionInFile(tree.getFunctionTrace(), traceWriter);
                operator.saveTrainingPredictionInFile(blockpredictionList.get(i), blockpredictionWriter);

                operator.saveTrainingPredictionInFile(originalStatementsList.get(i), originalStatementsWriter);
                operator.saveTrainingPredictionInFile(variableNameList.get(i), variableNamesWriter);

                operator.saveTrainingPredictionInFile((tempTree.getLines(tempTree.getRoot()) + 1)+"",linesWriter);

                operator.saveTrainingPredictionInFile( testCaseTraceList.get(i),testCaseTraceWriter);

                operator.saveTrainingPredictionInFile(blockOriginalStatementsList.get(i),blockOriginalStatementWriter);
                //}
            } catch (Exception e) {
                System.err.println(e.getMessage());
            } catch (Error e){
                System.err.println(e.getMessage());
            }
        }
    }

    /**
     * Construct training tree from the input code tree
     * Return the LinkedList with 2 list:
     * 1. The training trees with some continuous holes;
     * 2. The predication holes correspondent to the training trees;
     * */
    public LinkedList<LinkedList> getConstructTrainingData(CodeTree completeTree){
        // Init
        trees = new LinkedList<>();
        predications = new LinkedList<>();
        classnames = new LinkedList<>();
        parents = new LinkedList<>();
        holesizes = new LinkedList<>();
        blockpredictions = new LinkedList<>();
        originalStatements = new LinkedList<>();
        variableNames = new LinkedList<>();
        blockstatements = new LinkedList<>();

        testcaseTraces = new LinkedList<>();
        testcasePool = new HashMap<>(); // -- for test case
        sourceLines = new HashMap<>();

        // Construct
        int count = completeTree.getTotalNumber();
        for (int i = 2; i <= count; i++) {// construct from the second node
            construct(completeTree, i);
        }

        // Return
        LinkedList<LinkedList> result = new LinkedList<>();
        result.addLast(trees);
        result.addLast(predications);
        result.addLast(classnames);
        result.addLast(parents);
        result.addLast(holesizes);
        result.addLast(blockpredictions);
        result.addLast(originalStatements);
        result.addLast(variableNames);

        result.addLast(testcaseTraces);
        result.addLast(blockstatements);

        return result;
    }

    /**
     * @param completeTree: training data from this tree
     * @param serialNumber: the first node to be remove is with this serial number
     * */
    private void construct(CodeTree completeTree, int serialNumber){
        // Make the replica of the completeTree to avoid destroying.
        CodeTree tree = copyCodeTree(completeTree);
        operator.setSerialNumberofEachNode(tree);

        // Initialize the blockprediction
        blockprediction  = "";
        blockstatement = "";

        // The source code
        trace = Util.getTrace(completeTree);
        sourceLines = Util.getSourceLines(trace);

        if(tree.getTotalNumber() >= serialNumber){
            boolean isSuccess = true;
            List predicts = null;
            if(tree.getTotalNumber() == 1){// only root node.
                isSuccess = false;
            }
            else if((predicts = remove(tree,serialNumber)) == null
                    ){
                isSuccess = false;
            }
            if(isSuccess){
                String testcase = Util.convert2SourceCode(sourceLines);
                String predict = (String) predicts.get(0);

                if((!testcasePool.containsValue(testcase)) || isSpecialLabel(predict)) { // real, success when testcase pool do not contains the test case
                    count++;
                    if (isDebug) {
                        System.out.println(count + ". " + predicts);
                    }
                    predications.addLast(predict);
                    trees.addLast(tree);
                    classnames.addLast((String) predicts.get(1));
                    TreeNode generationNode = (TreeNode) predicts.get(2);
                    parents.addLast(generationNode);

                    // Add hole node
                    addHole(tree, generationNode);

                    // set as 1.
                    holesizes.addLast("1");
                    blockpredictions.addLast(blockprediction.trim());
                    if (isDebug) {
                        //System.out.println(blockprediction);
                        //displayTree(tree, true, blockprediction);
                    }

                    blockstatements.addLast(blockstatement.trim());
                    blockstatement = "";

                    // statements
                    originalStatements.addLast((String) predicts.get(3));
                    // variable names
                    variableNames.addLast((String) predicts.get(4));

                    // extract the function
                    String methodinfo = tree.getMethodInfo();
                    int b = Integer.parseInt(methodinfo.split(" ")[0]);
                    int e = Integer.parseInt(methodinfo.split(" ")[1]);
                    String functionblock = Util.convert2SourceCode(sourceLines, b, e);

                    // construct the testcase: testcase = sourceinfo + functionblock + }
                    String sourceinfo = tree.getSourceInfo();
                    testcase = sourceinfo + functionblock + "}";

                    // test case save
                    String testcaseTrace = saveTestCase(testcase, count, trace, b + "to" + e, "PreContext");
                    testcasePool.put(testcaseTrace,testcase);
                    // test case trace
                    testcaseTraces.addLast(testcaseTrace);
                }
            }
        }
    }

    private void addHole(CodeTree tree, TreeNode generationNode) {
        TreeNode hole = new TreeNode();
        hole.setClassName("hole");
        hole.setCompleteClassName("hole");
        hole.setMethodName("");
        hole.setCompleteMethodName("");
        hole.setAddMethodName(false);
        tree.addNode(tree.getTreeNode(generationNode.getSerialNumber()), hole);
        operator.setSerialNumberofEachNode(tree);
    }

    /**
     * Remove the node with serialNumber in the code tree
     * Remove it, including its children(except the remain node, e.g control)
     * Find the first ancestor node that under control node, remove the next brothers of this ancestor
     * Return its complete method declaration
     * */
    private LinkedList<Object> remove(CodeTree tree, int serialNumber){
        LinkedList<Object> result = new LinkedList<>();

        TreeNode node =  tree.getTreeNode(serialNumber);
        if(node == null || node.isCondition()){// Not exist
            return null;
        }
        TreeNode parent = node.getParentNode();
        String label = node.getCompleteMethodDeclaration();
        String className = node.getCompleteClassName();

        String statement = node.getStatement();
        List<String> previousVariableNames = node.getPreviousVariableNames();
        String variablename = previousVariableNames.size()>0?previousVariableNames.get(0):"";
        for (int i = 1; i < previousVariableNames.size(); i++) {
            variablename += " " + previousVariableNames.get(i);
        }

        // record predictions of this node
        blockprediction += " " + node2String(node);
        constructTestCase(node);

        // Find the first ancestor node that under control node
        TreeNode ancestor = findAncestorUnderControl(node);

        // Remove the next brother of this ancestor
        TreeNode elderAncestor;
        while(ancestor != null) {
            elderAncestor = ancestor.getParentNode();
            if (elderAncestor != null) {
                int index = elderAncestor.getChildNodes().indexOf(ancestor);
                int nextbrothernum = elderAncestor.getChildNodes().size() - (index + 1);
                for (int i = 1; i <= nextbrothernum; i++) {
                    blockprediction += " " + node2String(elderAncestor.getChildNodes().get(index + 1));
                    constructTestCase(elderAncestor.getChildNodes().get(index + 1));
                    elderAncestor.getChildNodes().remove(index + 1);
                }
            }
            ancestor = findAncestorUnderControl(elderAncestor);
        }

        // Remove the node
        if (parent != null) {
            parent.getChildNodes().remove(node);
        }

        // Reorder nodes
        operator.setSerialNumberofEachNode(tree);
        result.addLast(label);
        result.addLast(filterSigns(className));

        // Reassign hole parent
        if(isReassignHoleParent(label)){
            parent = node.getHoleParentNode();
        }

        if(parent == null) {
            parent = new TreeNode();
            parent.setSerialNumber(0);// the root do not have a parent
        }
        result.addLast(parent);
        result.addLast(statement);
        result.addLast(variablename);

        return result;
    }

    private TreeNode findAncestorUnderControl(TreeNode node) {
        TreeNode parent = null;
        TreeNode tmpnode = node;
        while((parent = tmpnode.getParentNode())!= null){
            if(parent.isControl()){
                return tmpnode;
            }
            tmpnode = parent;
        }
        return null;
    }

    // DFS node to string
    private String node2String(TreeNode treeNode) {
        String buff = treeNode.getCompleteMethodDeclaration();
        for (int i = 0; i < treeNode.getChildNodes().size(); i++) {
            buff += " " + node2String(treeNode.getChildNodes().get(i));
        }
        return buff;
    }

    // for testing ...
    private void displayTree(CodeTree codeTree, boolean isCompleteFlag,String title) {
        TreeView treeView = new TreeView();
        treeView.convertCodeTree(codeTree, isCompleteFlag);
        new DisplayTreeView(treeView.getTree(),title);
    }

    // filter the signs
    private String filterSigns(String string){
        for (String sign: filterSigns) {
            string = string.replace(sign, "");
        }
        return string;
    }

    // Reassign hole parent.
    private boolean isReassignHoleParent(String label){
        for(String l : holeParentAdjust){
            if(l.equals(label)){
                return true;
            }
        }
        return false;
    }

    private CodeTree copyCodeTree(CodeTree completeTree) {
        try {
            ByteArrayOutputStream bos = new ByteArrayOutputStream();
            ObjectOutputStream oos = new ObjectOutputStream(bos);
            oos.writeObject(completeTree);
            ByteArrayInputStream bis = new ByteArrayInputStream(bos.toByteArray());
            ObjectInputStream ois = new ObjectInputStream(bis);
            CodeTree tree = (CodeTree) ois.readObject();
            return tree;
        } catch (IOException | ClassNotFoundException e) {
            e.printStackTrace();
        }
        return null;
    }

    private String saveTestCase(String testcase, int count, String trace, String lines, String flag) {
        String path = testcaseDir + trace.replace("/", "_")+"_"+lines+"_"+count+"_"+flag+".java";
        try {
            BufferedWriter bufferedWriter = new BufferedWriter(new FileWriter(path));
            bufferedWriter.write(testcase);
            bufferedWriter.flush();
            bufferedWriter.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return path;
    }

    private boolean isSpecialLabel(String label){
        List<String> l = Arrays.stream(specialLabel).collect(Collectors.toList());
        return l.contains(label);
    }

    // Replace the begin~end lines as holes to construct testcase, including all its children.
    private void constructTestCase(TreeNode node) {

        String scinfo = node.getInfo();
        if(scinfo == null || node.isCondition()){
            return;
        }

        int begin = Util.getBeginLine(node);
        int end = Util.getEndLine(node);
        String stmt = Util.getStmt(node);

        if((!node.isControl()) && begin==end) {
            blockstatement = blockstatement + " " + Util.replaceStmt(begin, stmt, sourceLines," /*hole*/ ");
        }
        else {
            blockstatement = blockstatement + " " + Util.replaceLine2Line(begin, end, sourceLines," /*hole*/ ");
        }
        if(isDebug) {
            System.out.println("replacing " + node.getCompleteMethodDeclaration());
            System.out.println("b: " + begin + "; e: " + end);
            System.out.println(stmt);
            System.out.println();
            System.out.println(Util.convert2SourceCode(sourceLines));
        }
        List<TreeNode> children = node.getChildNodes();
        for (int i = 0; i < children.size(); i++) {
            constructTestCase(children.get(i));
        }
    }
}
